package com.jboltai.capability.aitext;

import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.jboltai.capability.Templates;
import com.jboltai.capability.functioncall.FuncCallResult;
import com.jboltai.capability.message.AIIntention;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventRunMode;
import com.jboltai.event.EventState;
import com.jboltai.event.ai.aifunccall.AIFuncCallEvent;
import com.jboltai.event.ai.aiintention.AIIntentionEvent;
import com.jboltai.event.ai.aiknowledge.AIKnowledgeEvent;
import com.jboltai.event.ai.aitext.AITextEvent;
import com.jboltai.event.ai.aitext.text2json.Text2JsonEvent;
import com.jboltai.event.ai.aitext.text2sql.SqlType;
import com.jboltai.event.ai.aitext.text2sql.Text2SqlEvent;
import com.jboltai.event.ai.rewrite.AIRewritePromptEvent;
import com.jboltai.resource.functioncall.FunctionResource;
import com.jboltai.resource.functioncall.FunctionResourceCenter;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Objects;

/**
 * ai文本类事件处理器
 */
public class AITextProcessor {

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    private AITextProcessor() {}

    public static AITextProcessor INSTANCE = new AITextProcessor();

    /**
     * 处理ai text 相关事件
     * @param event
     * @param mode
     */
    public void process(AITextEvent event, EventRunMode mode) {
        if (event.setState(EventState.PROCESSING)) {
            switch (event.getType()) {
                case AI_KNOWLEDGE:
                    processKnowledgeEvent((AIKnowledgeEvent)event, mode);
                    break;
                case AI_TEXT_TO_SQL:
                    processSQLEvent((Text2SqlEvent)event, mode);
                    break;
                case AI_TEXT_TO_JSON:
                    processJSONEvent((Text2JsonEvent)event, mode);
                    break;
                case AI_FUNC_CALL:
                    processFunctionCallEvent((AIFuncCallEvent)event, mode);
                    break;
                case AI_INTENTION:
                    processIntentionEvent((AIIntentionEvent)event, mode);
                    break;
                case AI_REWRITE_PROMPT:
                    processRewritePrompt((AIRewritePromptEvent)event, mode);
                    break;
            }
        }


    }

    /**
     * 处理意图识别事件
     * @param event
     * @param mode
     */
    private void processIntentionEvent(AIIntentionEvent event, EventRunMode mode) {

        event.getAiChatEvent().setAttr("intentions", event.getIntentions()).prompt(Templates.INTENTION_RECOGNIZE_TPL, event.getIntentionJSON(), event.getPrompt())
                .setStream(true)
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                })
                .onComplete((e) -> {
                    if (e.getState() != EventState.SUCCESS) return;
                    AIMessage result = e.getResult();
                    List<AIIntention> intentions =  e.getAttr("intentions");
                    AIIntention match = null;
                    String key = result.getContent().trim();
                    for (AIIntention intention : intentions) {
                        if (key.equals(intention.getKey()) || key.contains(intention.getKey())) {
                            match = intention;
                            break;
                        }
                    }
                    event.setResult(match == null ? AIIntention.NO_MATCH : match);
                    event.success(event.getResult());
                    event.setState(EventState.SUCCESS);
                    event.complete();
                }).onFail((e, error) -> {
                    event.fail(error.getErrorType(),error.getMsg(), error.getException());
                    LOGGER.error("IntentionEvent[{}]处理失败:{}", event.getId(), e);
                });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }

    }

    private void processKnowledgeEvent(AIKnowledgeEvent event, EventRunMode mode) {
        event.getAiChatEvent()
                .prompt(event.generateFinalPrompt())
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                })
                .onSuccess((e, result) -> {
                    event.success(result);
                }).onFail((e, error) -> {
                    event.fail(error.getErrorType(),error.getMsg(), error.getException());
                    LOGGER.error("KnowledgeEvent[{}]处理失败:{}", event.getId(), e);
                }).onComplete(e -> {
                    if (e.isSuccess()) {
                        event.setResult(e.getResult());
                        event.setState(EventState.SUCCESS);
                        //等待大模型，全部处理成功后，判断结果是否是无匹配内容
                        if (Objects.equals(e.getResult().getContent(), event.getNoMatchReply())) {
                            //未匹配到结果
                            if (event.getNoMatchHandler() != null) {
                                event.getNoMatchHandler().accept(event);
                            }
                        } else {
                            event.setMatchSuccess(true);
                        }
                        event.complete();
                    }
                });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }
    }

    /**
     * 处理function call 类型事件
     * @param event
     * @param mode
     */
    private void processFunctionCallEvent(AIFuncCallEvent event, EventRunMode mode) {
        event.getAiChatEvent().setMaxTokens(event.getMaxTokens())
                .prompt(Templates.FUNCTION_CALL_TPL, event.getFunctionResourceJSON(), event.getPrompt())
                .setStream(true)
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                }).onComplete((e) -> {
                    if (e.getState() != EventState.SUCCESS) return;
                    AIMessage result = e.getResult();
                    String content = result.getContent().trim();
                    LOGGER.debug("content:{}", content);
                    boolean isJson = JSONUtil.isTypeJSON(content);
                    if (!isJson) {
                        if (content.contains("{") && content.contains("}")) {
                            int lastIndex = content.lastIndexOf("function");
                            if (lastIndex > -1) {
                                content = content.substring(content.indexOf("{", (lastIndex-20) < 0?0:(lastIndex-20)), content.lastIndexOf("}") + 1);
                                isJson = JSONUtil.isTypeJSON(content);
                            }
                        }
                    }

                    FuncCallResult callResult = null;
                    if (isJson) {
                        callResult =  FuncCallResult.fromJson(content);
                    }
                    if (callResult == null || StrUtil.isBlank(callResult.getFunction())) {
                        event.fail(EventErrorType.NO_FUNCTION_CALL_MATCH);
                    } else {
                        try {
                            final String matchFunctionName = callResult.getFunction();
                            FunctionResource matchFunction = event.getFunctionResourceIds()
                                    .stream()
                                    .map(FunctionResourceCenter::getById)
                                    .filter(r -> Objects.equals(r.getName(), matchFunctionName))
                                    .findFirst().orElse(null);
                            if (matchFunction == null) {
                                LOGGER.error("匹配到的function[{}]已从注册中心取消注册", matchFunctionName);
                                event.fail(EventErrorType.NO_FUNCTION_CALL_MATCH);
                            } else {
                                callResult.setFunctionId(matchFunction.getId());
                                event.success(callResult);
                                event.setResult(callResult);
                                event.setState(EventState.SUCCESS);
                            }

                        } finally {
                            event.complete();
                        }
                    }



                }).onFail((e, error) -> {

                    event.fail(error.getErrorType(),error.getMsg(), error.getException());
                    LOGGER.error("FunctionCallEvent[{}]处理失败:{}", event.getId(), e);
                });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }
    }
    /**
     * 处理function call 类型事件
     * @param event
     * @param mode
     */
    private void processRewritePrompt(AIRewritePromptEvent event, EventRunMode mode) {
        event.getAiChatEvent().setMaxTokens(event.getMaxTokens())
                .prompt(Templates.PROMPT_REWRITE_TPL,event.getFormattedContextMsg(), event.getPrompt())
                .setStream(true)
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                }).onComplete((e) -> {
                    if (e.getState() != EventState.SUCCESS) return;
                    AIMessage result = e.getResult();
                    String content = result.getContent().trim();
                    LOGGER.debug("重写Prompt结果:{}", content);

                    if (StrUtil.isBlank(content) || content.contains("完整提问") == false) {
                        event.fail(EventErrorType.RESULT_INVALID);
                    } else {
                        try {
                            int index = content.indexOf("完整提问:");
                            if (index == -1) {
                                index = content.indexOf("完整提问") + 4;
                            } else {
                                index += 5;
                            }
                            String newPrompt = content.substring(index);
                            event.success(newPrompt);
                            event.setResult(newPrompt);
                            event.setState(EventState.SUCCESS);
                        } finally {
                            event.complete();
                        }
                    }



                }).onFail((e, error) -> {

                    event.fail(error.getErrorType(),error.getMsg(), error.getException());
                    LOGGER.error("AIRewritePromptEvent[{}]处理失败:{}", event.getId(), e);
                });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }
    }




    private String getJSONTemplate(Text2JsonEvent event) {

        StringBuilder template = new StringBuilder("请你充当一个JSON格式数据生成器、非结构化数据转JSON格式数据转换器，" +
                "将我提供的[原始数据]按照[JSON数据格式说明]进行转换，生成并输出符合这个JSON格式要求的数据，" +
                "请一步一步仔细思考后执行并精准回答，你不要做任何假设，不要输出任何思考过程，你最终只输出解析到的和转换后的JSON数据内容，不要输出任何跟JSON无关的内容，" +
                "如果无法转换成JSON就仅返回一个指定字符串:null：\n");
        if (StrUtil.isNotBlank(event.getBackground())) {
            template.append("背景知识如下：```").append(event.getBackground()).append("```。\n");
        }
        if (StrUtil.isNotBlank(event.getJsonDesc())) {
            template.append("转换后的JSON数据格式说明如下：```").append(event.getJsonDesc()).append("```。\n");
        }
        if (StrUtil.isNotBlank(event.getOriData())) {
            template.append("原始数据如下：```").append(event.getOriData()).append("```。\n");
        }
        if (StrUtil.isNotBlank(event.getPrompt())) {
            if (StrUtil.isBlank(event.getOriData())) {
                template.append("原始数据和");
            }
            template.append("其它转换生成要求如下：```").append(event.getPrompt()).append("```。\n");
        }

        return template.toString();
    }

    /**
     * 处理text转json类型事件
     * @param event
     * @param mode
     */
    private void processJSONEvent(Text2JsonEvent event, EventRunMode mode) {
        event.getAiChatEvent()
                .prompt(getJSONTemplate(event))
                .setStream(true)
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                }).onComplete((e) -> {
                    if (e.getState() != EventState.SUCCESS) return;
                    AIMessage result = e.getResult();
                    String content = result.getContent().trim();
                    if (content.startsWith("```json")) {
                        content = content.substring(7, content.length() - 3);
                    } else if (content.startsWith("```")) {
                        content = content.substring(3, content.length() - 3);
                    }
                    LOGGER.debug("content:{}", content);
                    if (!JSONUtil.isTypeJSON(content)) {
                        event.fail(EventErrorType.NOT_JSON);
                        return;
                    }

                    try {
                        event.success(content);
                        event.setResult(content);
                        event.setState(EventState.SUCCESS);
                    } finally {
                        event.complete();
                    }


                }).onFail((e, error) -> {

                    event.fail(error.getErrorType(),error.getMsg(), error.getException());
                    LOGGER.error("Text2JsonEvent[{}]处理失败:{}", event.getId(), e);
                });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }
    }


    /**
     * 处理text转sql类型事件
     * @param event
     * @param mode
     */
    private void processSQLEvent(Text2SqlEvent event, EventRunMode mode) {
        event.getAiChatEvent()
                .prompt(Templates.TEXT_TO_SQL,event.getBackground(), event.getPrompt(),event.getDBType().toString())
                .onThinking((e, think, status) -> {
                    if (event.getThinkHandler() != null) {
                        event.getThinkHandler().accept(event, think, status);
                    }
                })
                .setStream(true).onComplete((e) -> {
                    if (e.getState() != EventState.SUCCESS) return;
                    AIMessage result = e.getResult();
                    String content = result.getContent().trim();
                    if (content.startsWith("```sql")) {
                        content = content.substring(6, content.length() - 3);
                    } else if (content.startsWith("```")) {
                        content = content.substring(3, content.length() - 3);
                    }
                    content = content.trim();
                    LOGGER.debug("content:{}", content);
                    if (!content.contains(" ")) {
                        event.fail(EventErrorType.NOT_SQL);
                        return;
                    }
                    SqlType sqlType = SqlType.get(content.substring(0, content.indexOf(" ")));
                    if (sqlType == null) {
                        event.fail(EventErrorType.NOT_SQL);
                        return;
                    }

                    try {
                        CCJSqlParserUtil.parse(content);
                    } catch (JSQLParserException exception) {
                        event.fail(EventErrorType.NOT_SQL);
                        return;
                    }

                    if (event.getLimitSqlType() != null && event.getLimitSqlType().length > 0) {
                        //检查sqltype，假设获取的sql类似 select * from
                        if (!ArrayUtil.contains(event.getLimitSqlType(), sqlType)) {
                            event.fail(EventErrorType.SQL_TYPE_NOT_MATCH);
                            return;
                        }
                    }

                    event.success(content);
                    event.setResult(content);
                    event.setState(EventState.SUCCESS);
                    event.complete();

        }).onFail((e, error) -> {

            event.fail(error.getErrorType(),error.getMsg(), error.getException());
            LOGGER.error("Text2SqlEvent[{}]处理失败:{}", event.getId(), e);
        });
        switch (mode) {
            case PUBLISH:
                event.getAiChatEvent().publish();
                break;
            case DISPATCH:
                event.getAiChatEvent().dispatch();
                break;
        }
    }


}
